package main

import (
	"fmt"
	"os"
	"path"
	"sort"
	"sync/atomic"
	"time"

	"github.com/tecbot/gorocksdb"
)

type RocksdbStore struct {
	tableDef       TableDefine
	partition      Partition
	lastReadTime   atomic.Value
	lastWriteTime  atomic.Value
	lastOpenTime   atomic.Value
	storePath      string
	db             *gorocksdb.DB
	wo             *gorocksdb.WriteOptions
	ro             *gorocksdb.ReadOptions
	opt            *gorocksdb.Options
	datcfh, idxcfh *gorocksdb.ColumnFamilyHandle
}

var CFNames = []string{"default", "idx"}

func NewRocksdbStore(tableDef TableDefine, partition Partition, basePath string) (*RocksdbStore, error) {
	if tableDef.DtLength != len(partition.dt) {
		return nil, fmt.Errorf("dt length not match")
	}

	opt := gorocksdb.NewDefaultOptions()
	rateLimiter := gorocksdb.NewRateLimiter(102400, 100*1000, 10)
	opt.SetRateLimiter(rateLimiter)
	opt.SetCreateIfMissing(true)
	opt.SetCreateIfMissingColumnFamilies(true)

	ro := gorocksdb.NewDefaultReadOptions()
	wo := gorocksdb.NewDefaultWriteOptions()
	return &RocksdbStore{
		wo:        wo,
		ro:        ro,
		opt:       opt,
		tableDef:  tableDef,
		partition: partition,
		storePath: path.Join(basePath, tableDef.Path(), partition.StoreName()),
	}, nil
}

func (tbs *RocksdbStore) GetTableDef() TableDefine {
	return tbs.tableDef
}
func (tbs *RocksdbStore) GetPartition() Partition {
	return tbs.partition
}

func (tbs *RocksdbStore) Write(rows []*Row) error {
	wb := gorocksdb.NewWriteBatch()
	defer wb.Destroy()
	for _, row := range rows {
		wb.PutCF(tbs.datcfh, []byte(row.Key), []byte(row.Value))
		indexes, err := tbs.tableDef.BuildIndex(row)
		if err != nil {
			return err
		}
		for _, idx := range indexes {
			wb.PutCF(tbs.datcfh, []byte(idx), []byte{})
		}
		b, err := row.MarshallIndex()
		if err != nil {
			return err
		}
		wb.PutCF(tbs.idxcfh, []byte(row.Key), b)
	}
	if err := tbs.db.Write(tbs.wo, wb); err != nil {
		return err
	}
	tbs.lastWriteTime.Store(time.Now().Unix())
	return nil
}

func (tbs *RocksdbStore) LastActiveTime() (int64, int64, int64) {
	var o, r, w int64
	if tbs.lastOpenTime.Load() != nil {
		o = tbs.lastOpenTime.Load().(int64)
	}
	if tbs.lastReadTime.Load() != nil {
		r = tbs.lastReadTime.Load().(int64)
	}
	if tbs.lastWriteTime.Load() != nil {
		w = tbs.lastWriteTime.Load().(int64)
	}
	return o, r, w
}

func (tbs *RocksdbStore) Delete(k string) error {
	v, err := tbs.db.GetCF(tbs.ro, tbs.idxcfh, []byte(k))
	defer v.Free()
	if err != nil {
		return err
	}
	r := &Row{Key: k}
	if err := r.UnMarshallIndex(v.Data()); err != nil {
		return err
	}
	idxkeys, err := tbs.tableDef.BuildIndex(r)
	if err != nil {
		return err
	}
	wb := gorocksdb.NewWriteBatch()
	defer wb.Destroy()
	wb.DeleteCF(tbs.datcfh, []byte(k))
	wb.DeleteCF(tbs.idxcfh, []byte(k))
	for _, key := range idxkeys {
		wb.DeleteCF(tbs.datcfh, []byte(key))
	}
	if err := tbs.db.Write(tbs.wo, wb); err != nil {
		return err
	}
	tbs.lastWriteTime.Store(time.Now().Unix())
	return nil
}

func (tbs *RocksdbStore) Query(k string) (string, error) {
	v, err := tbs.db.GetCF(tbs.ro, tbs.datcfh, []byte(k))
	defer v.Free()
	if err != nil {
		return "", err
	}
	tbs.lastReadTime.Store(time.Now().Unix())
	return string(v.Data()), nil
}

func (tbs *RocksdbStore) QueryKeys(condition QueryCondition) ([]string, error) {
	var keys [][]string
	if condition.keyPrefix != "" {
		keys = append(keys, tbs.scan([]byte(condition.keyPrefix), false))
	}
	pres, err := tbs.tableDef.BuildQueryIndexPrefixes(condition)
	if err != nil {
		return nil, err
	}
	for _, p := range pres {
		k := tbs.scan([]byte(p), true)
		keys = append(keys, k)
	}
	for i, _ := range keys {
		if i != 0 {
			sort.Strings(keys[i])
		}
	}
	tbs.lastReadTime.Store(time.Now().Unix())
	return Keys(keys), nil
}

func (tbs *RocksdbStore) QueryValues(condition QueryCondition) (map[string]string, error) {
	ret := make(map[string]string)
	keys, err := tbs.QueryKeys(condition)
	if err != nil {
		return nil, err
	}
	for _, key := range keys {
		value, err := tbs.Query(key)
		if err != nil {
			return nil, err
		}
		ret[key] = value
	}
	tbs.lastReadTime.Store(time.Now().Unix())
	return ret, nil
}

func (tbs *RocksdbStore) scan(prefix []byte, isIndex bool) (ret []string) {
	iter := tbs.db.NewIteratorCF(tbs.ro, tbs.datcfh)
	defer iter.Close()
	for iter.Seek(prefix); iter.ValidForPrefix(prefix); iter.Next() {
		raw := iter.Key().Data()
		if isIndex {
			if key, err := tbs.tableDef.ExtractKey(raw); err == nil && len(key) > 0 {
				ret = append(ret, string(key))
			}
			continue
		}
		if len(raw) > 0 {
			ret = append(ret, string(raw))
		}
	}
	return ret
}

func (tbs *RocksdbStore) Open() error {
	if _, err := os.Stat(tbs.storePath); err != nil {
		if !os.IsNotExist(err) {
			return err
		}
		err = os.MkdirAll(tbs.storePath, 0644)
		if err != nil {
			return err
		}

	}
	//db, err := gorocksdb.OpenDb(tbs.opt, tbs.storePath)
	db, cfh, err := gorocksdb.OpenDbColumnFamilies(tbs.opt, tbs.storePath, CFNames, []*gorocksdb.Options{tbs.opt, tbs.opt})
	if err != nil {
		return err
	}
	tbs.db = db
	tbs.datcfh = cfh[0]
	tbs.idxcfh = cfh[1]
	tbs.lastOpenTime.Store(time.Now().Unix())
	return nil
}

func (tbs *RocksdbStore) Close() error {
	if tbs.db != nil {
		tbs.datcfh.Destroy()
		tbs.idxcfh.Destroy()
		tbs.db.Close()
		tbs.db = nil
		tbs.datcfh = nil
		tbs.idxcfh = nil
	}
	return nil
}

func (tbs *RocksdbStore) IsActive() bool {
	return tbs.db != nil && tbs.datcfh != nil && tbs.idxcfh != nil
}
